module poly_1d_mod
  use recon_type_mod
  use recon_math_mod
  use poly_utils_mod

  implicit none
  
contains

  subroutine poly_1d_init(this, nd, sw, swx, swy, swz, xc, yc, zc, dx, dy, dz, ic, jc, kc, is, ie, js, je, ks, ke, mask, id)

    class(recon_type), intent(inout) :: this
    integer , intent(in), optional :: nd
    integer , intent(in), optional :: sw
    integer , intent(in), optional :: swx
    integer , intent(in), optional :: swy
    integer , intent(in), optional :: swz
    real(16), intent(in), optional :: xc(:)
    real(16), intent(in), optional :: yc(:)
    real(16), intent(in), optional :: zc(:)
    real(16), intent(in), optional :: dx
    real(16), intent(in), optional :: dy
    real(16), intent(in), optional :: dz
    integer , intent(in), optional :: ic
    integer , intent(in), optional :: jc
    integer , intent(in), optional :: kc
    integer , intent(in), optional :: is
    integer , intent(in), optional :: ie
    integer , intent(in), optional :: js
    integer , intent(in), optional :: je
    integer , intent(in), optional :: ks
    integer , intent(in), optional :: ke
    integer , intent(in), optional :: mask(:,:,:) ! Cell mask
    integer , intent(in), optional :: id

    integer i, j, k

    call this%clear()

    if(nd/=1)then
      stop 'poly_1d_init support only nd==1'
    endif
    
    this%sw  = sw
    this%swx = sw
    this%swy = 1
    this%swz = 1
    this%nd  = nd
    this%rw  = ( sw - 1 ) / 2

    if(present(ic))then
      this%ic = ic
    else
      this%ic = this%rw + 1
    endif

    this%dx = dx
    if(present(dy))then
      this%dy = dy
    else
      this%dy = 1
    endif
    if(present(dz))then
      this%dz = dz
    else
      this%dz = 1
    endif

    if(present(is))then; this%is = is; else; this%is = 1       ; endif
    if(present(ie))then; this%ie = ie; else; this%ie = this%swx; endif
    if(present(js))then; this%js = js; else; this%js = 1       ; endif
    if(present(je))then; this%je = je; else; this%je = 1       ; endif
    if(present(ks))then; this%ks = ks; else; this%ks = 1       ; endif
    if(present(ke))then; this%ke = ke; else; this%ke = 1       ; endif

    allocate(this%cell_mask(this%sw,1,1))
    allocate(this%poly_mask(this%sw,1,1))

    if (present(mask)) then
      this%cell_mask = mask
    else
      this%cell_mask = 1
    end if

    this%nc = count(this%cell_mask == 1)

    this%rwL = this%rw
    this%rwR = this%rw
    do i = 0, -this%rw, -1
      if( this%cell_mask(this%ic+i,1,1) == 1 )then
        this%rwL = min(abs(i), this%rw)
      endif
    enddo
    do i = 0, this%rw
      if( this%cell_mask(this%ic+i,1,1) == 1 )then
        this%rwR = min(abs(i), this%rw)
      endif
    enddo

    this%poly_mask = this%cell_mask

    allocate(this%xc(this%ic-this%rwL:this%ic+this%rwR))
    if (present(xc)) then
      this%xc = xc
    else
      ! Set coordinates of cells on the large stencil with origin at center.
      do i = -this%rwL, this%rwR
        this%xc(this%ic+i) = i
      end do
    end if
    
    this%get_jab => poly_1d_get_jacobian

    this%initialized = .true.

  end subroutine poly_1d_init

  subroutine poly_1d_calc_recon_matrix(this, ierr)

    class(recon_type), intent(inout) :: this
    integer, intent(out) :: ierr

    ! Local double double arrays for preserving precision.
    real(16), allocatable, dimension(:,:) :: A, iA
    integer i, j, k, ipt, n, iterm, dxn

    ierr = 0

    if (allocated(this%iA            )) deallocate(this%iA            )
    if (allocated(this%poly          )) deallocate(this%poly          )
    if (allocated(this%poly_r16      )) deallocate(this%poly_r16      )
    if (allocated(this%recon_mtx     )) deallocate(this%recon_mtx     )
    if (allocated(this%recon_mtx_r16 )) deallocate(this%recon_mtx_r16 )
    if (allocated(this%dpoly         )) deallocate(this%dpoly         )
    if (allocated(this%dpoly_r16     )) deallocate(this%dpoly_r16     )
    if (allocated(this%drecon_mtx    )) deallocate(this%drecon_mtx    )

    allocate(this%iA            (this%nc ,this%nc)) ! (nterm, nc)
    allocate(this%poly          (this%npt,this%nc))
    allocate(this%poly_r16      (this%npt,this%nc))
    allocate(this%recon_mtx     (this%npt,this%nc))
    allocate(this%recon_mtx_r16 (this%npt,this%nc))
    allocate(this%dpoly         (this%npt,this%nc,0:this%nc-1,0:0,0:0))
    allocate(this%dpoly_r16     (this%npt,this%nc,0:this%nc-1,0:0,0:0))
    allocate(this%drecon_mtx    (this%npt,this%nc,0:this%nc-1,0:0,0:0))

    allocate( A(this%nc,this%nc))
    allocate(iA(this%nc,this%nc))

    ! Set the poly for each evaluation point.
    ! Select monomials according to mask.
    do ipt = 1, this%npt
      iterm = 1
      do i = 1, this%nc
        call calc_monomial(this%x(ipt), iterm - 1, this%poly_r16(ipt,iterm))
        do dxn = 0, this%nc - 1
          call calc_deriv_monomial(this%x(ipt), iterm - 1, dxn, this%nc, this%dpoly_r16(ipt,iterm,dxn,0,0))
          this%dpoly_r16(ipt,iterm,dxn,0,0) = this%dpoly_r16(ipt,iterm,dxn,0,0) / this%dx**dxn
        end do
        iterm = iterm + 1
      end do
    end do
    
    ! Calculate inverse of integral coefficient matrix.
    call calc_tensor_product_poly_integral_matrix(this%nc, this%xc, A)
    call inverse_matrix(A, iA, ierr)
    if (ierr /= 0) then
      print*,'Failed to calculate inverse A in poly_1d_calc_recon_matrix'
      deallocate(A, iA)
      return
    end if
    
    this%iA    = iA
    this%poly  = this%poly_r16
    this%dpoly = this%dpoly_r16
    
    this%recon_mtx_r16 = matmul(this%poly_r16, iA)
    this%recon_mtx = this%recon_mtx_r16

    do dxn = 0, this%nc - 1
      this%drecon_mtx(:,:,dxn,0,0) = matmul(this%dpoly_r16(:,:,dxn,0,0), iA)
    end do
    
    deallocate(A, iA)

  end subroutine poly_1d_calc_recon_matrix
  
  subroutine poly_1d_reconstruct(this, fi, fo, TCI, ierr)

    class(recon_type), intent(inout) :: this
    real(8), intent(in ) :: fi(:,:,:) ! Cell averaged function values
    real(8), intent(out) :: fo(:)     ! Reconstructed function values on evaluation points
    logical, intent(in ), optional :: TCI ! Trouble Cell Indicator(TCI), 1 for existing TC(use WENO), 0 for no TC(use poly)
    integer, intent(out), optional :: ierr

#ifdef DEBUG
    if(present(ierr))then
      ierr = 0
      if (size(fi) /= this%nc) then
        if (present(ierr)) ierr = 1
        return
      end if
      if (size(fo) /= this%npt) then
        if (present(ierr)) ierr = 1
        return
      end if
    endif
#endif

    fo = matmul(this%recon_mtx, pack(fi, .true.))
    
  end subroutine poly_1d_reconstruct
  
  subroutine poly_1d_reconstruct_deriv(this, fi, fx, fy, fz, ierr)

    class(recon_type), intent(inout) :: this
    real(8), intent(in ) :: fi(:,:,:) ! Cell averaged function values
    real(8), intent(out), optional :: fx(:)     ! Reconstructed x derivative values on evaluation points
    real(8), intent(out), optional :: fy(:)     ! Reconstructed x derivative values on evaluation points
    real(8), intent(out), optional :: fz(:)     ! Reconstructed x derivative values on evaluation points
    integer, intent(out), optional :: ierr

#ifdef DEBUG
    if(present(ierr))then
      ierr = 0
      if (size(fi) /= this%nc) then
        if (present(ierr)) ierr = 1
        return
      end if
      if (size(fo) /= this%npt) then
        if (present(ierr)) ierr = 1
        return
      end if
    endif
#endif

    fx = matmul(this%drecon_mtx(:,:,1,0,0), pack(fi, .true.))
    
  end subroutine poly_1d_reconstruct_deriv
  
  subroutine poly_1d_get_jacobian(this, dqLdq, dqRdq, q)
    class(recon_type), intent(in) :: this
    real(8), dimension(this%nc), intent(inout) :: dqLdq
    real(8), dimension(this%nc), intent(inout) :: dqRdq
    real(8), dimension(this%nc), intent(in   ), optional :: q    
    
    dqLdq = this%recon_mtx(1,:)
    dqRdq = this%recon_mtx(2,:)
    
  end subroutine poly_1d_get_jacobian

end module poly_1d_mod
